/*
 Copyright 2023 YaoZinan   zinan@outlook.com   an-archi.cc

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

      https://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */

#include <arpa/inet.h>
#include <errno.h>
#include <event2/buffer.h>
#include <event2/bufferevent.h>
#include <event2/event.h>
#include <event2/listener.h>
#include <event2/util.h>
#include <map>
#include <netinet/in.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sys/socket.h>

#include "trace.h"
#include "tunne.h"

#define MSG_PAYLOAD_REF( _type_, _msg_ ) ( *( _type_* )( ( char* )( &_msg_ ) + sizeof( MsgHeader ) ) )

namespace lrp {
int TunnelMgmt::_s_id = 1;
TunnelMgmt::TunnelMgmt( const Station& sta_ )
    : _sta( sta_ )
    , _ev_base( nullptr ) {
    init();
}

int TunnelMgmt::gen_id() {
    return _s_id++;
}

int TunnelMgmt::init() {
    _ev_base = event_base_new();

    if ( _sta == Station::Client ) create_timer();

    return _ev_base ? 0 : -1;
}

void TunnelMgmt::run() {
    if ( !ev_base() ) return;

    event_base_dispatch( ev_base() );
}

TunnelMgmt::~TunnelMgmt() {
    shutdown_spots();
    event_base_free( ev_base() );
    event_free( _timer );
}

void TunnelMgmt::shutdown_spots() {
    // todo crash
    //  for ( auto& s : _srv_spots ) {
    //      evconnlistener_free( s->lis() );
    //  }
}

void TunnelMgmt::close( int chan_ ) {
    if ( chan_ == 0 ) {
        shutdown_spots();
    }
    _tunnels.erase( chan_ );
    delete tunnel_by_id( chan_ );
}

void TunnelMgmt::listener_cb( struct evconnlistener* listener,
                              evutil_socket_t        fd,
                              struct sockaddr*       sa,
                              int                    socklen,
                              void*                  user_data ) {
    Spot* spot = ( Spot* )user_data;

    struct bufferevent* bev = bufferevent_socket_new( spot->mgmt().ev_base(), fd, BEV_OPT_CLOSE_ON_FREE );
    if ( !bev ) {
        TRACE( "Error constructing bufferevent!" );
        // todo 这里有一个问题，不能因为一个socket连接出错就关闭所有的
        if ( spot->is_main() ) {
            // event_base_loopbreak( self->_ev_base );
        }
        else {
            // todo
        }
        return;
    }

    // todo
    evutil_make_socket_nonblocking( fd );
    struct sockaddr_in* sin = ( struct sockaddr_in* )sa;

    TRACE( "New connection from %s:%d.\n", inet_ntoa( sin->sin_addr ), ntohs( sin->sin_port ) );

    Tunnel* t = nullptr;
    if ( spot->is_main() )
        t = new SrvTunnel( spot->mgmt(), bev, 0 );
    else
        t = new SideTunnel( spot->mgmt(), bev, gen_id() );

    t->bind_addr( inet_ntoa( sin->sin_addr ), ntohs( sin->sin_port ) );

    bufferevent_setcb( bev, read_cb, nullptr, event_cb, t );
    bufferevent_enable( bev, EV_WRITE | EV_READ );
}

int TunnelMgmt::create_srv_spot( uint16_t port_, bool main_ ) {
    if ( !ev_base() ) {
        TRACE( "ev_base is null\n" );
        return -1;
    }
    TRACE( "srv pot,port=%d, main=%d\n", port_, main_ );

    auto spot = std::make_unique<Spot>( *this, nullptr, main_ );

    struct sockaddr_in sin = { 0 };
    sin.sin_family         = AF_INET;
    sin.sin_port           = htons( port_ );

    auto listener = evconnlistener_new_bind( ev_base(),
                                             listener_cb,
                                             spot.get(),
                                             LEV_OPT_REUSEABLE | LEV_OPT_CLOSE_ON_FREE,
                                             -1,
                                             ( struct sockaddr* )&sin,
                                             sizeof( sin ) );

    if ( !listener ) {
        TRACE( "Could not create a listener! erro=%s\n", strerror( errno ) );
        return 1;
    }

    spot->set_listener( listener );

    TRACE( "Listening on port %d.\n", port_ );
    // evutil_make_socket_nonblocking( listener );
    _srv_spots.push_back( spot.release() );

    return 0;
}

void TunnelMgmt::tunnel_unimped( int chan_ ) {
    auto t = tunnel_by_id( chan_ );
    if ( !t ) return;
    t->set_through();
}

void TunnelMgmt::event_cb( struct bufferevent* bev_, short events_, void* user_data_ ) {
    auto tunnel = reinterpret_cast<Tunnel*>( user_data_ );

    if ( events_ & BEV_EVENT_EOF ) {
        TRACE( "Connection closed.\n" );
        tunnel->disconnected( false );
    }
    else if ( events_ & BEV_EVENT_ERROR ) {
        TRACE( "Got an error on the connection: %s\n", strerror( errno ) ); /*XXX win32*/
        // return;
        tunnel->disconnected( true );
    }
    else if ( events_ & BEV_EVENT_CONNECTED ) {
        // 连接成功时走这里，并且要客户端第一次触发读事件后连接才真正建立起来
        tunnel->connected();
        TRACE( "connect success\n" );
    }
}

void TunnelMgmt::read_cb( struct bufferevent* bev_, void* user_data_ ) {
    reinterpret_cast<Tunnel*>( user_data_ )->recv();
}

void TunnelMgmt::add_pending( const Tunnel* t_ ) {
    _pending_cli.push_back( const_cast<Tunnel*>( t_ ) );
}

//
int TunnelMgmt::forward( int index_, const void* data_, int length_ ) {
    auto t = tunnel_by_id( index_ );
    if ( !t ) return -1;

    return length_ == t->send( data_, length_ );
}

Tunnel* TunnelMgmt::tunnel_by_id( int id_ ) {
    return _tunnels.count( id_ ) > 0 ? _tunnels[ id_ ] : nullptr;
}

void TunnelMgmt::remove_pending( const Tunnel* t_ ) {
    _pending_cli.remove( const_cast<Tunnel*>( t_ ) );
}

void TunnelMgmt::timer_cb( evutil_socket_t fd, short event, void* arg ) {
    TunnelMgmt* self = ( TunnelMgmt* )arg;
    TRACE( "*cli to conn=%d\n", ( int )self->_pending_cli.size() );

    for ( auto& t : self->_pending_cli ) {
        bufferevent_socket_connect( t->bev(), ( struct sockaddr* )&t->addr(), sizeof( t->addr() ) );
    }
}

void TunnelMgmt::create_timer() {
    _timer = event_new( _ev_base, -1, EV_TIMEOUT | EV_PERSIST, &TunnelMgmt::timer_cb, this );
    if ( !_timer ) {
        TRACE( "event_new failed\n" );
        exit( -1 );
    }

    struct timeval ti {
        0, 500000  // 500ms.
    };

    event_add( _timer, &ti );
}

Tunnel* TunnelMgmt::create_cli_tunnel( const char* ip_, uint16_t port_, int chan_ ) {
    struct bufferevent* bev = bufferevent_socket_new( ev_base(), -1, BEV_OPT_CLOSE_ON_FREE );
    if ( !bev ) {
        TRACE( "socket init failed\n" );
        return nullptr;
    }

    Tunnel* t = 0 == chan_ ? ( Tunnel* )new CliTunnel( *this, bev, chan_ ) : new SideTunnel( *this, bev, chan_ );

    bufferevent_setcb( bev, read_cb, nullptr, event_cb, t );
    bufferevent_enable( bev, EV_READ | EV_WRITE );

    // 每个client tunnel都有一个服务器地址，每个serv tunnel也有一个本地
    TRACE( "cli tunnel par: ip=%s,port=%d,chan=%d\n", ip_, port_, chan_ );
    t->bind_addr( ip_, port_ );

    _pending_cli.push_back( t );

    if ( bufferevent_socket_connect( bev, ( struct sockaddr* )&t->addr(), sizeof( t->addr() ) ) < 0 ) {
        TRACE( "connect failed\n" );
        return nullptr;
    }

    return t;
}

void Tunnel::set_through() {
    _get_through = true;
}

void Tunnel::close() {
    bufferevent_free( _bev );
}

void Tunnel::connected() {
    TRACE( "base-- connected\n" );
    if ( mgmt().station() == Station::Client ) {
        _mgmt.remove_pending( this );
    }
}

void Tunnel::disconnected( bool abnormal_ ) {
    TRACE( "base-- disconnected\n" );

    if ( 0 != id() && mgmt().station() == Station::Server ) {
        close();
    }

    if ( mgmt().station() == Station::Client ) {
        mgmt().add_pending( this );
    }
}

Tunnel::Tunnel( TunnelMgmt& mgmt_, struct bufferevent* bev_, int id_ )
    : _mgmt( mgmt_ )
    , _get_through( false )
    , _id( id_ )
    , _received( 0 )
    , _length( sizeof( MsgHeader ) )
    , _bev( bev_ )
    , _cache( nullptr )
    , _cc_length( 0 ) {
}

void Tunnel::assign( int id_, struct bufferevent* bev_ ) {
    _id  = id_;
    _bev = bev_;
}

Tunnel::~Tunnel() {
    close();
}

void Tunnel::bind_addr( const char* ip_, uint16_t port_ ) {
    _sin.sin_addr.s_addr = inet_addr( ip_ );
    _sin.sin_family      = AF_INET;
    _sin.sin_port        = htons( port_ );
}

int Tunnel::send( const void* data_, size_t sz_ ) {
    TRACE( "get through %d\n", _get_through );
#if 0  // todo 只需要在cli那一头缓冲？服务端是否需要
    if ( !_get_through ) {
        if ( sz_ + _cc_length > kTunnelCacheSize )
            return -1;

        if ( !_cache ) {
            _cache     = std::make_unique<char[]>( kTunnelCacheSize );
            _cc_length = 0;
        }

        memcpy( _cache.get() + _cc_length, data_, sz_ );
        _cc_length += sz_;
        return sz_;
    }

    if ( _cc_length > 0 ) {
        bufferevent_write( _bev, _cache.get(), _cc_length );
        // std::swap( _cache, nullptr );
        _cache.reset( 0 );
        _cc_length = 0;
    }
#endif
    return ( int )bufferevent_write( _bev, data_, sz_ );
}

void CliTunnel::on_data( int chan_, void* data_, int length_ ) {
    TRACE( "cli data\n" );
}

void SrvTunnel::on_data( int chan_, void* data_, int length_ ) {
    TRACE( "srv data\n" );
}

void CliTunnel::on_conn( int chan_, uint16_t port_ ) {
    auto tunnel = _mgmt.create_cli_tunnel( "127.0.0.1", port_, chan_ );
    if ( !tunnel ) confirm( false );
}

void SrvTunnel::on_confirm( int chan_, const MsgConnConfirm& confirm_ ) {
    if ( !confirm_.success ) {
        _mgmt.close( chan_ );
    }
    else {
        _mgmt.tunnel_unimped( chan_ );
    }
}

void SrvTunnel::on_config( const MsgConfig& cfg_ ) {
    for ( int i = 0; i < cfg_.count; ++i ) {
        _mgmt.create_srv_spot( cfg_.cfg[ i ].r_port, false );
    }
}

void SrvTunnel::on_heartbeat() {
    TRACE( "srv->heartbeat\n" );
}

void CliTunnel::on_heartbeat() {
    TRACE( "cli->heartbeat\n" );
}

void CliTunnel::confirm( bool valid_ ) {
    MsgHeader msg = { 0 };
    msg.id        = ( int )MsgId::MSG_CONN_CONFIRM;
    msg.chan      = _id;
    msg.length    = sizeof( MsgConnConfirm );
    memcpy( _message.data(), &msg, sizeof( msg ) );

    MsgConnConfirm confirm{ valid_ };
    memcpy( _message.data() + sizeof( MsgHeader ), &confirm, sizeof( confirm ) );

    send( _message.data(), sizeof( MsgHeader ) + sizeof( MsgConnConfirm ) );
}

void CliTunnel::config() {
    TRACE( "cli->config\n" );

    MsgHeader& msg = reinterpret_cast<MsgHeader&>( *_message.data() );
    msg.id         = ( int )MsgId::MSG_CONFIG;
    msg.chan       = _id;

    MsgConfig& cfg = reinterpret_cast<MsgConfig&>( *( _message.data() + sizeof( MsgHeader ) ) );

    cfg.count = _cfgs.size();
    for ( int i = 0; i < cfg.count; ++i ) {
        cfg.cfg[ i ] = _cfgs[ i ];
    }

    msg.length = sizeof( cfg.count ) + cfg.count * sizeof( TunnelCfg );

    send( _message.data(), sizeof( MsgHeader ) + msg.length );
}

void MainTunnel::close() {
    // 关闭其他所有通道
    _mgmt.shutdown_spots();
    Tunnel::close();
}

void MainTunnel::on_unknown( const MsgHeader& msg_ ) {
    TRACE( "unknown msg\n" );
}

// todo 服务器和客户端要处理的消息是不一样的
void MainTunnel::msg_received( const MsgHeader& msg_ ) {
    on_unknown( msg_ );
}

void MainTunnel::recv() {
    auto sz = evbuffer_get_length( bufferevent_get_input( _bev ) );
    TRACE( "main recv %lu,exp=%lu,recv=%lu\n", sz, _length, _received );

    while ( sz > 0 ) {
        size_t r = std::min( sz, _length - _received );
        bufferevent_read( _bev, _message.data() + _received, r );
        _received = r;

        if ( r + _received >= _length ) {
            auto payload = reinterpret_cast<const MsgHeader&>( *_message.data() ).length;

            if ( payload == 0 || _length > sizeof( MsgHeader ) ) {
                msg_received( reinterpret_cast<MsgHeader&>( *( _message.data() ) ) );
            }
            else {
                assert( _length == sizeof( MsgHeader ) );
                _length += payload;
            }
        }

        sz -= r;
    }
}

void SrvTunnel::msg_received( const MsgHeader& msg_ ) {
    switch ( ( MsgId )msg_.id ) {
    default: return MainTunnel::msg_received( msg_ );
    case MsgId::MSG_DATA:          // -> lrpc,->lrps
        return on_data( msg_.chan, ( char* )&msg_ + sizeof( MsgHeader ), msg_.length );
    case MsgId::MSG_CONN_CONFIRM:  // ->lrps
        return on_confirm( msg_.chan, MSG_PAYLOAD_REF( MsgConnConfirm, msg_ ) );
    case MsgId::MSG_CONFIG:        // ->lrps
        return on_config( MSG_PAYLOAD_REF( MsgConfig, msg_ ) );
    case MsgId::MSG_HEARTBEAT:
        return on_heartbeat();
    }
}

void CliTunnel::connected() {
    config();

    // if ( 0 == _id ) {
    //     TRACE( "send config\n" );
    // }
    // else {
    //     TRACE( "config proxy conn\n" );
    //     confirm( true );
    // }
}

void CliTunnel::msg_received( const MsgHeader& msg_ ) {
    switch ( ( MsgId )msg_.id ) {
    default: return MainTunnel::msg_received( msg_ );
    case MsgId::MSG_DATA:  // -> lrpc,->lrps
        return on_data( msg_.chan, ( char* )&msg_ + sizeof( MsgHeader ), msg_.length );
    case MsgId::MSG_CONN:  // -> lrpc
        return;            // todo on_conn( msg_ );
    case MsgId::MSG_HEARTBEAT:
        return on_heartbeat();
    }
}

void SideTunnel::recv() {
    auto sz = evbuffer_get_length( bufferevent_get_input( _bev ) );
    while ( sz > 0 ) {
        auto r = std::min( kTunnelBufferSize, ( int )sz );

        MsgHeader header;

        header.id     = ( int )MsgId::MSG_DATA;
        header.length = sz;
        header.chan   = id();

        _mgmt.forward( 0, &header, sizeof( header ) );
        bufferevent_read( _bev, _message.data(), r );
        auto rc = process( _message.data(), r );
        _mgmt.forward( 0, _message.data(), rc );
        sz -= r;
    }
    // todo
    //  bufferevent_write_buffer( bev, bufferevent_get_input( bev ) );
}

}  // namespace lrp